package com.dd.stack.es.service.task.data.gene.person;

import com.dd.stack.es.ESApp;
import com.dd.stack.es.pojo.es.PersonESDO;
import net.datafaker.Faker;
import org.springframework.boot.SpringApplication;
import org.springframework.context.ConfigurableApplicationContext;
import org.springframework.data.elasticsearch.core.ElasticsearchOperations;
import org.springframework.data.elasticsearch.core.query.IndexQuery;
import org.springframework.data.elasticsearch.core.query.IndexQueryBuilder;

import java.time.LocalDateTime;
import java.util.ArrayList;
import java.util.List;
import java.util.UUID;
import java.util.stream.Collectors;

/**
 * @Author liuxianmeng
 * @CreateTime 2025/4/21 0:24
 * @Description 测试数据生成
 */
public class PersonGenerator {
    public static void main(String[] args) {
        // 初始化Spring上下文（需要主配置类）
        ConfigurableApplicationContext context = SpringApplication.run(ESApp.class, args);

        try {
            // 获取Elasticsearch操作模板
            ElasticsearchOperations operations = context.getBean(ElasticsearchOperations.class);

            // 创建生成器实例并执行
            PersonGenerator generator = new PersonGenerator(operations);
            generator.generateAndSavePersons(10000);

            System.out.println("数据生成完成！");
        } finally {
            // 确保关闭上下文
            context.close();
        }
    }

    private static final Faker FAKER = new Faker();
    private static final int BATCH_SIZE = 500;

    private final ElasticsearchOperations operations;

    public PersonGenerator(ElasticsearchOperations operations) {
        this.operations = operations;
    }

    public void generateAndSavePersons(int count) {
        System.out.printf("开始生成 %d 条测试数据...%n", count);

        List<PersonESDO> batch = new ArrayList<>(BATCH_SIZE);
        for (int i = 1; i <= count; i++) {
            batch.add(createFakePerson());

            if (i % BATCH_SIZE == 0) {
                saveBatch(batch);
                batch.clear();
                System.out.printf("已生成 %d/%d 条（%.1f%%）%n",
                        i, count, (i * 100.0)/count);
            }
        }

        if (!batch.isEmpty()) {
            saveBatch(batch);
        }
        System.out.println("所有数据已生成完成！");
    }

    private PersonESDO createFakePerson() {
        PersonESDO person = new PersonESDO();
        person.setId(UUID.randomUUID().toString());
        person.setName(FAKER.name().fullName());
        person.setAge(FAKER.number().numberBetween(1, 100));
        person.setHeight((float) FAKER.number().randomDouble(2, 1, 2));
        person.setBirthDate(FAKER.date().birthday(18, 65).toLocalDateTime().toLocalDate());
        person.setAddress(createFakeAddress());
        person.setCreatedAt(LocalDateTime.now());
        return person;
    }

    private void saveBatch(List<PersonESDO> batch) {
        try {
            List<IndexQuery> queries = batch.stream()
                    .map(p -> new IndexQueryBuilder().withObject(p).build())
                    .collect(Collectors.toList());

            operations.bulkIndex(queries, PersonESDO.class);
        } catch (Exception e) {
            System.err.println("批量插入失败: " + e.getMessage());
            e.printStackTrace();
        }
    }

    private PersonESDO.Address createFakeAddress() {
        PersonESDO.Address address = new PersonESDO.Address();
        address.setStreet(FAKER.address().streetAddress());
        address.setZipCode(Integer.valueOf(FAKER.address().zipCode().split("-")[0]));
        address.setCoordinates(createFakeCoordinates());
        return address;
    }

    private PersonESDO.Coordinates createFakeCoordinates() {
        PersonESDO.Coordinates coordinates = new PersonESDO.Coordinates();
        coordinates.setLatitude(Float.valueOf(FAKER.address().latitude()));
        coordinates.setLongitude(Float.valueOf(FAKER.address().longitude()));
        coordinates.setDescription(FAKER.lorem().sentence());
        coordinates.setTestNum(FAKER.number().numberBetween(1000, 9999));
        return coordinates;
    }
}
